#include <iostream>
#include <eigen3/Eigen/Eigen>
#include <EKF.h>
#include <cmath>
#include <random>
#include <matplotlibcpp.h>

using namespace Eigen;
using namespace std;
namespace plt = matplotlibcpp;

void Init(Vector4d x0, Vector4d P0, Matrix2Xd &z, Matrix4Xd &x_re)
{
    double T = 0.1, k_x = 0.01, k_y = 0.05, g = 9.8;
    Eigen::Vector2d v;
    v << 8, 0.001;
    mt19937 gen_xx, gen_xy, gen_zr, gen_za;
    normal_distribution<double> dis_xx(0.0, sqrt(P0(1)));
    normal_distribution<double> dis_xy(0.0, sqrt(P0(3)));
    normal_distribution<double> dis_zr(0.0, sqrt(v(0)));
    normal_distribution<double> dis_za(0.0, sqrt(v(1)));
    int count = z.cols();
    x_re.col(0)(0) = x0(0);
    x_re.col(0)(1) = x0(1);
    x_re.col(0)(2) = x0(2);
    x_re.col(0)(3) = x0(3);
    z.col(0)(0) = sqrt(pow(x_re.col(0)(0), 2) + pow(x_re.col(0)(2), 2)) + dis_zr(gen_zr);
    z.col(0)(1) = atan(x_re.col(0)(0) / x_re.col(0)(2)) + dis_za(gen_za);
    for (int i = 1; i < count; i++)
    {
        x_re.col(i)(0) = x_re.col(i-1)(0) + T * x_re.col(i-1)(1);
        x_re.col(i)(1) = x_re.col(i-1)(1) - k_x * T * pow(x_re.col(i-1)(1), 2) + dis_xx(gen_xx);
        x_re.col(i)(2) = x_re.col(i-1)(2) + T * x_re.col(i-1)(3);
        x_re.col(i)(3) = x_re.col(i-1)(3) + k_y * T * pow(x_re.col(i-1)(3), 2) - g * T + dis_xy(gen_xy);
        z.col(i)(0) = sqrt(pow(x_re.col(i)(0), 2) + pow(x_re.col(i)(2), 2)) + dis_zr(gen_zr);
        z.col(i)(1) = atan2(x_re.col(i)(0), x_re.col(i)(2)) + dis_za(gen_za);
    }
}

int main(int argc, char *argv[])
{
    double T_s = 15;
    double T = 0.1;
    int size = T_s / T;
    // 初始化
    Vector4d x0;
    Vector4d Pd;
    x0 << 0, 50, 500, 0;
    Pd << 0, 0.3, 0, 0.3;
    Matrix4d P0 = Pd.asDiagonal();
    EKF plane(x0, P0);
    Matrix2Xd z = Matrix2Xd::Zero(2, size);
    Matrix4Xd x_re = Matrix4Xd::Zero(4, size);
    Matrix4Xd x = Matrix4Xd::Zero(4, size);
    Init(x0, Pd, z, x_re);
    plane.Pushz_k(z);
    plane.EKFAlgo();
    plane.Getx_k(x);
    vector<double> x_EKF[] = {vector<double>(size), vector<double>(size)};
    vector<double> z_EKF[] = {vector<double>(size), vector<double>(size)};
    vector<double> x_real[] = {vector<double>(size), vector<double>(size)};
    int i = 0;
    for (; i < size; i++)
    {
        x_EKF[0][i] = x.col(i)(0);
        x_EKF[1][i] = x.col(i)(2);
        z_EKF[0][i] = z.col(i)(0) * sin(z.col(i)(1));
        z_EKF[1][i] = z.col(i)(0) * cos(z.col(i)(1));
        x_real[0][i] = x_re.col(i)(0);
        x_real[1][i] = x_re.col(i)(2);
    }
    plt::title("result");
    plt::xlabel("x");
    plt::ylabel("y");
    plt::named_plot("EKF estimated", x_EKF[0], x_EKF[1], "r-.");
    plt::named_plot("measurement", z_EKF[0], z_EKF[1], "b*");
    plt::named_plot("real", x_real[0], x_real[1], "g--");
    plt::legend();
    plt::save("./result.png");
    plt::show();
    return 0;
}